#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/12/1 10:54
# @Author  : huidong.bai
# @File    : ResultStorage.py
# @Software: PyCharm
# @Mail    : MasterBai2018@outlook.com

import pdb
import copy
from queue import Queue
from conftest import logging


class ResultStorage:
    def __init__(self, yaml_config):
        self.config = yaml_config
        self.result_storage = self._init_storage()

    def _init_storage(self):
        storage = {}
        queues = {}
        for i in range(4):
            storage[str(i)] = copy.deepcopy(queues)
            for result_type, fields in self.config.items():
                storage[str(i)][result_type] = {"Queue": Queue(), "Processed": False}
        return storage

    def set_finish_cond(self):
        for i in range(4):
            for result_type, fields in self.config.items():
                self.result_storage[str(i)][result_type]["Queue"].put(None)
    
    def clear_queue(self, channel, result_type):
        self.result_storage[channel][result_type]["Queue"].queue.clear()

    def checkout_queue_result(self, channel, msg_type):
        try:
            result = None
            processed = self.result_storage[channel][msg_type]["Processed"]
            if not processed:
                data = self.result_storage[channel][msg_type]["Queue"].get()
                if data is None:
                    self.result_storage[channel][msg_type]["Processed"] = True
                else:
                    self.result_storage[channel][msg_type]["Processed"] = False
                    result = data
            return result
        except Exception as e:
            logging.error(f"Get result storage queue: {channel}, {msg_type}, error: {e}")
            return None

    def collect_result(self, data):
        message_type = data.get("resultType")
        channel_id = str(data.get("channelId"))
        self.result_storage[channel_id][message_type]["Queue"].put(data)

    def reset(self):
        del self.result_storage
        self.result_storage = self._init_storage()
